#coding:utf8

import numpy as np
import theano
import theano.tensor as T
import matplotlib.pyplot as plt
from scipy import signal
import time
np.random.seed(123)

def generateData(M,T,dB,L):
    #QPSK信号
    TxS = np.sign(np.random.rand(M) * 2 - 1) + 1j*np.sign(np.random.rand(M) * 2 - 1) #30000
    #ch = np.random.randn(chL+1) + 1j * np.random.randn(chL+1)
    ch = [0.0410+0.0109j,0.0495+0.0123j,0.0672+0.0170j,0.0919+0.0235j,
     0.7920+0.1281j,0.3960+0.0871j,0.2715+0.0498j,0.2291+0.0414j,0.1287+0.0154j,
     0.1032+0.0119j]
    ch = ch / np.linalg.norm(ch)
    x = signal.fftconvolve(ch,TxS)[:M]   #信道卷积 x.shape = (30000,)
    #noise
    n=np.random.randn(1,M)+1j*np.random.randn(1,M);
    n=n/np.linalg.norm(n)*pow(10,(-dB/20))*np.linalg.norm(x);
    x = x + n
    x = x[0]
    K = M-L-1 #29987
    X = []
    for i in range(K):
        X.append(x[i+L+1:i:-1])
    X = np.array(X).T  # (13,29987)
    TxS = TxS[L:M-6]
    Y = X[:,5:]
    #Y为训练集（接收机收到的数据），TxS为target(发送序列), x为接收端接收到的数据(用来画图比较)
    return Y.T,TxS.T,x

class Layer(object):
    def __init__(self, inputs, in_size, out_size, activation_function=None):
        self.W = theano.shared(np.random.normal(0,1,(in_size, out_size)))
        self.Wx_plus_b = T.dot(inputs, self.W)
        self.activation_function = activation_function
        if activation_function:
            self.outputs = self.activation_function(self.Wx_plus_b)
        else:
            self.outputs = self.Wx_plus_b

dB = 25
smoothingLen = 12
#x接受数据，为一个复数序列，
#Y训练数据为smoothingLen列的复数矩阵 19990 * 13
#Tx为label，也是一个复数序列 19990 * 1
Y,Tx,x = generateData(30000,20000,dB,smoothingLen)
Y_real = np.real(Y)
Y_imag = np.imag(Y)
Y = np.hstack((Y_real,Y_imag)) #19990 * 26
Tx_real = np.real(Tx)
Tx_imag = np.imag(Tx)
Tx = np.vstack((Tx_real, Tx_imag)).T  #19990 * 2


# determine the inputs
z = T.dmatrix('z')
y = T.dmatrix('y')

# add layers
l1 = Layer(y, 2*smoothingLen+2, 2, None)  #26入2出

# loss function 误差函数
cost = T.mean(T.square(l1.outputs - z))

# compute the gradients
gW1 = T.grad(cost, l1.W)
# apply the gradient descent
learning_rate = 0.05
train = theano.function(
    inputs = [y,z],
    outputs = cost,
    updates = [(l1.W, l1.W - learning_rate * gW1)])

# predict
predict = theano.function(inputs=[y], outputs=l1.outputs)

# # train
start = time.time()
for i in range(301):
    err = train(Y, Tx)
    if i % 50 == 0:
        print i, time.time()-start, err

testY,testTx,testx = generateData(30000,20000,dB,smoothingLen)
testY = np.hstack((np.real(testY),np.imag(testY)))
predictValue = predict(testY)
print predictValue.shape

fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.scatter(predictValue[::2],predictValue[1::2])
plt.show()
